import json
import os
import copy
import math
import numpy as np
import re
import xml.dom.minidom
import xml.etree.ElementTree as ET

from pyproj import Proj


def parse_openscenario(file_dir: str):
    vehicles = {}
    type_dict = {}

    tree = ET.parse(file_dir)
    root = tree.getroot()
    # Entities = root.find('Entities')
    ScenarioObject_list = root.find('Entities').findall('ScenarioObject')
    for ScenarioObject in ScenarioObject_list:
        # print(ScenarioObject.attrib['name'])
        name = ScenarioObject.attrib['name']
        Vehicle = ScenarioObject.find('Vehicle')
        Pedestrian = ScenarioObject.find('Pedestrian')
        if Vehicle:
            vehicleCategory = Vehicle.attrib['vehicleCategory']
            if vehicleCategory == 'car':
                type_dict[name] = 1
            if vehicleCategory == 'bicycle':
                type_dict[name] = 6
        if Pedestrian:
            type_dict[name] = 7

    # 读取OpenScenario文件
    opens = xml.dom.minidom.parse(file_dir).documentElement

    """以下读取背景车相关信息，车辆编号从1号开始，轨迹信息记录在vehicle_traj中"""
    """新版场景中采用更general的定义方式，在初始时仅初始化主车，背景车的采用Event中的AddEntityAction和DeleteEntityAction"""
    act_list = opens.getElementsByTagName('Act')
    for act in act_list:
        Vertex_list = act.getElementsByTagName('Vertex')
        # print(act.getAttribute('name')[4:])
        id = act.getAttribute('name')[4:]
        # 记录OpenScenario中存在的位置、航向角、时间信息
        for Vertex in Vertex_list:
            t = float(Vertex.getAttribute('time'))
            x = float(Vertex.getElementsByTagName('WorldPosition')[0].getAttribute('x'))  # 记录横向位置
            y = float(Vertex.getElementsByTagName('WorldPosition')[0].getAttribute('y'))  # 记录纵向位置
            yaw = float(Vertex.getElementsByTagName('WorldPosition')[0].getAttribute('h'))  # 记录航向角
            yaw = (yaw + 2 * math.pi) if -math.pi <= yaw < 0 else yaw  # 航向角范围调整到(0, 2pi)
            v = 0
            if len(Vertex_list) > 1:
                v = np.sqrt((float(Vertex_list[1].getElementsByTagName('WorldPosition')[0].getAttribute('x')) - float(
                    Vertex_list[0].getElementsByTagName('WorldPosition')[0].getAttribute('x'))) ** 2 + (float(
                    Vertex_list[1].getElementsByTagName('WorldPosition')[0].getAttribute('y')) - float(
                    Vertex_list[0].getElementsByTagName('WorldPosition')[0].getAttribute('y'))) ** 2) / (
                            float(Vertex_list[1].getAttribute('time')) - float(Vertex_list[0].getAttribute('time')))
            if t not in vehicles.keys():
                vehicles[t] = [[id, x, y, yaw, v, type_dict[str(id)]]]
            else:
                vehicles[t].append([id, x, y, yaw, v, type_dict[str(id)]])

    return vehicles


def parse_json(file_dir: str):
    vehicles = {}
    type_dict = {}

    with open(file_dir, 'r', encoding='utf-8') as f:
        # 加载JSON数据到字典
        data = json.load(f)
        participantTrajectories = data['participantTrajectories']
        for participantTrajectory in participantTrajectories:
            model = participantTrajectory['model']
            trajectory = participantTrajectory['trajectory']
            name = participantTrajectory['name']
            if name == "A0" or participantTrajectory['id'] == "1":
                continue
            if model == 1:
                type_dict[name] = 1
            if model == 4:
                type_dict[name] = 7
            if model == 13:
                type_dict[name] = 6
            for traj in trajectory:
                position = traj['position']
                lon = position[0]
                lat = position[1]
                traj_time = float(traj['time'])
                p = Proj('+proj=tmerc +lon_0=108.90575652010739 +lat_0=34.37650478465651 +ellps=WGS84')
                x, y = p(lon, lat, inverse=False)
                if traj_time not in vehicles.keys():
                    vehicles[traj_time] = [[name, x, y, 0, 0, type_dict[str(name)]]]
                else:
                    vehicles[traj_time].append([name, x, y, 0, 0, type_dict[str(name)]])
    return vehicles


def parse_json_object(participantTrajectories: dict):
    vehicles = {}
    type_dict = {}

    for participantTrajectory in participantTrajectories:
        model = participantTrajectory['model']
        trajectory = participantTrajectory['trajectory']
        name = participantTrajectory['name']
        if name == "A0" or participantTrajectory['id'] == "1":
            continue
        if model == 1:
            type_dict[name] = 1
        if model == 4:
            type_dict[name] = 7
        if model == 13:
            type_dict[name] = 6
        for traj in trajectory:
            position = traj['position']
            lon = position[0]
            lat = position[1]
            traj_time = float(traj['time'])
            p = Proj('+proj=tmerc +lon_0=108.90575652010739 +lat_0=34.37650478465651 +ellps=WGS84')
            x, y = p(lon, lat, inverse=False)
            if traj_time not in vehicles.keys():
                vehicles[traj_time] = [[name, x, y, 0, 0, type_dict[str(name)]]]
            else:
                vehicles[traj_time].append([name, x, y, 0, 0, type_dict[str(name)]])
    return vehicles


if __name__ == '__main__':

    data = (parse_json(
        r"F:\辛英杰\project\cd_nine\cd_playoff_wanji_Process\net_file\output0715\R1.json"))
    for k, v in data.items():
        print(k, v)

